
import cv2
import time

from head_pose_est import (get_image_points, get_pose_estimation, get_euler_angle)

class HeadPose():
    """
        Detect the pose of head.
    """

    def __init__(self):
        super().__init__()

    @staticmethod
    def get(img):
        """
            Get the pose of head.
            Input:
                img, image contains a face.
            Return:
                head_pose, pose of head(eular angle)
        """
        start_time = time.time()
        # resize the image when it is oversize.
        size = img.shape
        if size[0] > 700:
            h = size[0] / 3
            w = size[1] / 3
            img = cv2.resize( img, (int( w ), int( h )), interpolation=cv2.INTER_CUBIC )
            size = img.shape
        # get the face feature points.
        ret, image_points = get_image_points(img)
        if ret != 0:
            print('get_image_points failed')
            return
        # get the face pose.
        ret, rotation_vector, translation_vector, camera_matrix, dist_coeffs = get_pose_estimation(size, image_points)
        if ret != True:
            print('get_pose_estimation failed')
            return
            
        used_time = time.time() - start_time
        print("used_time:{} sec".format(round(used_time, 3)))
        # convert rotation to euler_angle.
        ret, pitch, yaw, roll = get_euler_angle(rotation_vector)
        head_pose = [pitch, yaw, roll]
        euler_angle_str = 'Y:{}, X:{}, Z:{}'.format(pitch, yaw, roll)
        print(euler_angle_str)

        return head_pose


# test
if __name__ == "__main__":
    
    img = cv2.imread("test/test_1.jpg")
    is_front = HeadPose.get(img)
    print(is_front)
